import visdom
import numpy as np
import torch
import time


# 1. 连接测试函数
def test_visdom_connection(port=8097):
    try:
        vis = visdom.Visdom(port=port)
        return vis.check_connection()
    except Exception as e:
        print(f"连接失败: {str(e)}")
        return False


# 2. 可视化演示类
class VisdomDemo:
    def __init__(self, env_name='demo'):
        self.vis = visdom.Visdom(env=env_name)
        if not self.vis.check_connection():
            raise RuntimeError("Visdom服务器未启动！请先运行: python -m visdom.server")

    def text_demo(self):
        """文本展示"""
        self.vis.text('Hello Visdom!', win='text_win')

    def line_plot(self):
        """动态折线图"""
        win = self.vis.line(
            X=np.array([0]),
            Y=np.array([0]),
            opts=dict(title='动态正弦曲线', showlegend=True)
        )
        for i in range(1, 100):
            x = torch.FloatTensor([i / 10.])
            y = torch.sin(x)
            self.vis.line(
                X=x,
                Y=y,
                win=win,
                update='append',
                name='sin(x)'
            )
            time.sleep(0.1)

    def image_demo(self):
        """图像展示"""
        # 生成随机图片
        self.vis.images(
            torch.randn(16, 3, 64, 64),
            nrow=4,
            win='images_win',
            opts=dict(title='随机生成图片')
        )

    def scatter_plot(self):
        """散点图"""
        x = torch.randn(100, 2)
        y = (x[:, 0] > 0).long()
        self.vis.scatter(
            x,
            Y=y + 1,
            opts=dict(
                markersize=10,
                title='二分类散点图',
                legend=['Class 0', 'Class 1']
            ),
            win='scatter_win'
        )


# 3. 执行测试
if __name__ == "__main__":
    # 先检测连接
    if test_visdom_connection():
        print("Visdom连接成功！开始可视化演示...")

        demo = VisdomDemo(env_name='test_env')

        # 执行各可视化功能
        demo.text_demo()
        demo.scatter_plot()
        demo.image_demo()
        demo.line_plot()
    else:
        print("请先启动Visdom服务：python -m visdom.server")